#include "CustomHidController.h"
#include "RGBStates.h"

#include <cassert>
#include <iostream>

namespace TLAC::Input::RawInputCustomHidController
{
	void CustomHidControllerDevice::BeginTickMessages()
	{
		input_buttons_tapped_ = {};
		input_buttons_released_ = {};
	}

	bool CustomHidControllerDevice::OpenDevice(HANDLE handle)
	{
		if (!device_.OpenHandle(handle)) {
			return false;
		}

		// Find out which input usages are available in which reports
		inputs_by_report_id_.clear();

		auto MarkInputButton = [&](UCHAR report_id, USAGE usage, USHORT data_index) {
			if (usage >= BUTTON_USAGE_MIN && usage <= BUTTON_USAGE_MAX) {
				auto& button = inputs_by_report_id_[report_id].buttons[usage - BUTTON_USAGE_MIN];
				button.is_present = true;
				button.data_index = data_index;
			}
		};

		for (auto const& button_cap : device_.GetSpecificButtonCaps(HidP_Input, BUTTON_USAGE_PAGE, 0, 0)) {
			if (button_cap.IsRange) {
				USHORT range_length = button_cap.Range.UsageMax - button_cap.Range.UsageMin + 1;
				for (USHORT i = 0; i < range_length; ++i) {
					MarkInputButton(button_cap.ReportID, button_cap.Range.UsageMin + i, button_cap.Range.DataIndexMin + i);
				}
			}
			else {
				MarkInputButton(button_cap.ReportID, button_cap.NotRange.Usage, button_cap.NotRange.DataIndex);
			}
		}

		auto MarkInputSliderValue = [&](UCHAR report_id, USAGE usage, USHORT data_index, LONG logical_max) {
			if (usage >= TOUCH_SLIDER_USAGE_MIN && usage <= TOUCH_SLIDER_USAGE_MAX) {
				auto& value = inputs_by_report_id_[report_id].slider_values[usage - TOUCH_SLIDER_USAGE_MIN];
				value.is_present = true;
				value.data_index = data_index;
				value.logical_max = logical_max;
			}
		};

		/* Note: An HID Usage that has logical range 0-1 is considered a "button", while
		* a Usage with a larger range is a "value". They are handled completely differently
		* by Windows HID API.
		* A slider implementation that only transmits touch on/off will probably be treated
		* as buttons, while a slider that transmits a range of touch values will be treated
		* as values. Both needs to be handled.
		*/

		for (auto const& button_cap : device_.GetSpecificButtonCaps(HidP_Input, DIVA_USAGE_PAGE, 0, 0)) {
			if (button_cap.IsRange) {
				USHORT range_length = button_cap.Range.UsageMax - button_cap.Range.UsageMin + 1;
				for (USHORT i = 0; i < range_length; ++i) {
					MarkInputSliderValue(button_cap.ReportID, button_cap.Range.UsageMin + i, button_cap.Range.DataIndexMin + i, 1);
				}
			}
			else {
				MarkInputSliderValue(button_cap.ReportID, button_cap.NotRange.Usage, button_cap.NotRange.DataIndex, 1);
			}
		}

		for (auto const& value_cap : device_.GetSpecificValueCaps(HidP_Input, DIVA_USAGE_PAGE, 0, 0)) {
			if (value_cap.IsRange) {
				USHORT range_length = value_cap.Range.UsageMax - value_cap.Range.UsageMin + 1;
				for (USHORT i = 0; i < range_length; ++i) {
					MarkInputSliderValue(value_cap.ReportID, value_cap.Range.UsageMin + i, value_cap.Range.DataIndexMin + i, value_cap.LogicalMax);
				}
			}
			else {
				MarkInputSliderValue(value_cap.ReportID, value_cap.NotRange.Usage, value_cap.NotRange.DataIndex, value_cap.LogicalMax);
			}
		}

		inputs_available_ = {};

		for (auto& [report_id, record] : inputs_by_report_id_) {
			record.update_presents();
			for (std::size_t i = 0; i < inputs_available_.buttons.size(); ++i) {
				inputs_available_.buttons[i] |= record.buttons[i].is_present;
			}
			for (std::size_t i = 0; i < inputs_available_.values.size(); ++i) {
				inputs_available_.values[i] |= record.slider_values[i].is_present;
			}
			inputs_available_.has_any_button |= record.buttons_present;
			inputs_available_.has_any_value |= record.slider_present;
		}

		// Find out which output usages are available in which reports
		auto MarkOutputValue = [&](UCHAR report_id, USAGE usage, USHORT data_index, LONG logical_max) {
			UsageToDataIndex* value{ nullptr };
			if (usage == LED_GLOBAL_BRIGHTNESS_USAGE) {
				value = &outputs_by_report_id_[report_id].slider_global_brightness;
			}
			else if (usage >= LED_BRG_USAGE_MIN && usage <= LED_BRG_USAGE_MAX) {
				value = &outputs_by_report_id_[report_id].slider_leds[usage - LED_BRG_USAGE_MIN];
			}
			if (value) {
				value->is_present = true;
				value->data_index = data_index;
				value->logical_max = logical_max;
			}
		};

		for (auto const& value_cap : device_.GetSpecificValueCaps(HidP_Output, DIVA_USAGE_PAGE, 0, 0)) {
			if (value_cap.IsRange) {
				USHORT range_length = value_cap.Range.UsageMax - value_cap.Range.UsageMin + 1;
				for (USHORT i = 0; i < range_length; ++i) {
					MarkOutputValue(value_cap.ReportID, value_cap.Range.UsageMin + i, value_cap.Range.DataIndexMin + i, value_cap.LogicalMax);
				}
			}
			else {
				MarkOutputValue(value_cap.ReportID, value_cap.NotRange.Usage, value_cap.NotRange.DataIndex, value_cap.LogicalMax);
			}
		}

		auto MarkOutputButton = [&](UCHAR report_id, USAGE usage, USHORT data_index) {
			UsageToDataIndex* button{ nullptr };
			if (usage >= BUTTON_LIGHTS_USAGE_MIN && usage <= BUTTON_LIGHTS_USAGE_MAX) {
				button = &outputs_by_report_id_[report_id].button_lights[usage - BUTTON_LIGHTS_USAGE_MIN];
			}
			else if (usage >= PARTITION_LIGHTS_USAGE_MIN && usage <= PARTITION_LIGHTS_USAGE_MAX) {
				button = &outputs_by_report_id_[report_id].partition_lights[usage - PARTITION_LIGHTS_USAGE_MIN];
			}
			if (button) {
				button->is_present = true;
				button->data_index = data_index;
			}
		};

		for (auto const& button_cap : device_.GetSpecificButtonCaps(HidP_Output, DIVA_USAGE_PAGE, 0, 0)) {
			if (button_cap.IsRange) {
				USHORT range_length = button_cap.Range.UsageMax - button_cap.Range.UsageMin + 1;
				for (USHORT i = 0; i < range_length; ++i) {
					MarkOutputButton(button_cap.ReportID, button_cap.Range.UsageMin + i, button_cap.Range.DataIndexMin + i);
				}
			}
			else {
				MarkOutputButton(button_cap.ReportID, button_cap.NotRange.Usage, button_cap.NotRange.DataIndex);
			}
		}

		pending_output_writes_.clear();
		for (auto const& [report_id, record] : outputs_by_report_id_) {
			auto& output_report = pending_output_writes_[report_id];
			device_.InitializeReportForID(HidP_Output, report_id, output_report.buffer);
		}

		return true;
	}

	void CustomHidControllerDevice::HandleInput(RAWINPUT& input, std::vector<SliderValuesReport>& slider_reports)
	{
		assert(input.header.dwType == RIM_TYPEHID);
		for (std::size_t i = 0; i < input.data.hid.dwCount; ++i) {
			BYTE* report = input.data.hid.bRawData + (i * input.data.hid.dwSizeHid);
			const UCHAR report_id = report[0];
			auto it = inputs_by_report_id_.find(report_id);
			if (it == inputs_by_report_id_.end()) continue;
			const auto& inputs_in_report = it->second;

			const auto& data = device_.GetDataFromInputReport(reinterpret_cast<PCHAR>(report), input.data.hid.dwSizeHid);

			if (inputs_in_report.buttons_present) {
				for (std::size_t i = 0; i < num_input_buttons; ++i) {
					if (!inputs_in_report.buttons[i].is_present) {
						continue;
					}
					if (std::any_of(data.begin(), data.end(), [&](const HIDP_DATA& d) {
						return d.DataIndex == inputs_in_report.buttons[i].data_index;
					}))
					{
						// button is down
						if (!input_buttons_state_[i]) {
							input_buttons_tapped_[i] = true;
							input_buttons_state_[i] = true;
						}
					}
					else {
						// button is up
						if (input_buttons_state_[i]) {
							input_buttons_released_[i] = true;
							input_buttons_state_[i] = false;
						}
					}
				}
			}

			if (inputs_in_report.slider_present) {
				SliderValuesReport report_values{};
				for (std::size_t i = 0; i < num_input_slider_values; ++i) {
					if (!inputs_in_report.slider_values[i].is_present) continue;
					report_values.has_values[i] = true;
					auto it = std::find_if(data.begin(), data.end(), [&](const HIDP_DATA& d) {
						return d.DataIndex == inputs_in_report.slider_values[i].data_index;
					});
					ULONG raw_value;
					if (it != data.end()) {
						/* Either this usage is a "value", or this usage is a "button" and it's on.
						* Since HIDP_DATA has union{ULONG RawValue; BOOLEAN On;}, simply reading
						* RawValue should cover both cases.
						*/
						raw_value = it->RawValue;
					}
					else {
						/* This usage is a "button" and it's off.
						* Or, maybe *data* does not contain this usage at all for some reason,
						* but that would mean the API had gone wrong, so we ignore this possibility.
						*/
						raw_value = 0;
					}
					if (inputs_in_report.slider_values[i].logical_max == TOUCH_SLIDER_DEFAULT_MAX) {
						report_values.values[i] = raw_value;
					}
					else {
						report_values.values[i] = raw_value * TOUCH_SLIDER_DEFAULT_MAX / inputs_in_report.slider_values[i].logical_max;
					}
				}
				slider_reports.push_back(report_values);
			}
		}
	}

	void CustomHidControllerDevice::EndTickMessages()
	{
		if (outputs_by_report_id_.empty()) {
			return;
		}

		auto slider_array_ptr = RawInputDetails::RGBStates::GetSliderBRGArray();
		if (slider_array_ptr == nullptr) {
			return;
		}
		auto partition_lights_ptr = RawInputDetails::RGBStates::GetPartitionLightsRGBPackedByte();
		if (partition_lights_ptr == nullptr) {
			return;
		}

		auto slider_array = *slider_array_ptr;
		auto button_lights = RawInputDetails::RGBStates::GetButtonLightsPackedByte();
		auto partition_lights = *partition_lights_ptr;

		for (auto& [report_id, report] : pending_output_writes_) {
			if (report.IsPending()) {
				std::wcout << L"[TLAC] CustomHidControllerDevice::EndTickMessages(): output report " << int(report_id) << L" still pending" << std::endl;
				continue;
			}
			output_data_buffer_.clear();
			auto InsertDataButton = [&](USHORT data_index) {
				output_data_buffer_.push_back({});
				output_data_buffer_.back().DataIndex = data_index;
				output_data_buffer_.back().On = TRUE;
			};
			auto InsertDataValue = [&](USHORT data_index, ULONG raw_value) {
				output_data_buffer_.push_back({});
				output_data_buffer_.back().DataIndex = data_index;
				output_data_buffer_.back().RawValue = raw_value;
			};
			device_.InitializeReportForID(HidP_Output, report_id, report.buffer);
			auto const& record = outputs_by_report_id_.at(report_id);
			if (record.slider_global_brightness.is_present) {
				ULONG raw_value = slider_array[0];
				if (record.slider_global_brightness.logical_max != LED_GLOBAL_BRIGHTNESS_DEFAULT_MAX) {
					raw_value = raw_value * record.slider_global_brightness.logical_max / LED_GLOBAL_BRIGHTNESS_DEFAULT_MAX;
				}
				InsertDataValue(record.slider_global_brightness.data_index, raw_value);
			}
			for (std::size_t i = 0; i < num_output_led_brg_values; ++i) {
				if (record.slider_leds[i].is_present) {
					ULONG raw_value = slider_array[i + 1];
					if (record.slider_leds[i].logical_max != LED_BRG_DEFAULT_MAX) {
						raw_value = raw_value * record.slider_leds[i].logical_max / LED_BRG_DEFAULT_MAX;
					}
					InsertDataValue(record.slider_leds[i].data_index, raw_value);
				}
			}
			for (std::size_t i = 0; i < num_button_lights; ++i) {
				if (record.button_lights[i].is_present) {
					if (button_lights & (1 << i)) {
						InsertDataButton(record.button_lights[i].data_index);
					}
				}
			}
			for (std::size_t i = 0; i < num_partition_lights; ++i) {
				if (record.partition_lights[i].is_present) {
					if (partition_lights & (1 << i)) {
						InsertDataButton(record.partition_lights[i].data_index);
					}
				}
			}
			if (device_.SetDataForOutputReport(output_data_buffer_, report.buffer)) {
				report.Write(device_.GetFileHandle());
			}
		}
	}


	void CustomHidControllerHandler::AddDevicesToRegister(std::vector<RAWINPUTDEVICE>& devices) const
	{
		devices.push_back({ TLC_USAGE_PAGE, TLC_USAGE, RIDEV_DEVNOTIFY });
	}

	void CustomHidControllerHandler::BeginTickMessages()
	{
		for (auto& [handle, device] : devices_) {
			device.BeginTickMessages();
		}
		input_slider_reports_.clear();
	}

	void CustomHidControllerHandler::HandleDeviceArrival(HANDLE handle, const RID_DEVICE_INFO& device_info)
	{
		if (device_info.dwType != RIM_TYPEHID) return;

		if (device_info.hid.usUsagePage == TLC_USAGE_PAGE && device_info.hid.usUsage == TLC_USAGE) {
			if (devices_[handle].OpenDevice(handle)) {
				printf("[TLAC] Raw Input: Custom HID Controller: opened device\n");
			}
			else {
				devices_.erase(handle);
			}
			UpdateInputsAvailable();
		}
	}

	void CustomHidControllerHandler::HandleDeviceRemoval(HANDLE handle)
	{
		if (devices_.erase(handle) > 0) {
			printf("[TLAC] Raw Input: Custom HID Controller: closed device\n");
			UpdateInputsAvailable();
		}
	}

	void CustomHidControllerHandler::HandleInput(RAWINPUT& input)
	{
		if (input.header.dwType != RIM_TYPEHID) return;

		auto device_it = devices_.find(input.header.hDevice);
		if (device_it != devices_.end()) {
			device_it->second.HandleInput(input, input_slider_reports_);
		}
	}

	void CustomHidControllerHandler::EndTickMessages()
	{
		for (auto& [handle, device] : devices_) {
			device.EndTickMessages();
		}
		CollectInputStates();
	}

	void CustomHidControllerHandler::UpdateInputsAvailable()
	{
		inputs_available_ = {};
		for (auto& [handle, device] : devices_) {
			for (std::size_t i = 0; i < inputs_available_.buttons.size(); ++i) {
				inputs_available_.buttons[i] |= device.inputs_available_.buttons[i];
			}
			for (std::size_t i = 0; i < inputs_available_.values.size(); ++i) {
				inputs_available_.values[i] |= device.inputs_available_.values[i];
			}
			inputs_available_.has_any_button |= device.inputs_available_.has_any_button;
			inputs_available_.has_any_value |= device.inputs_available_.has_any_value;
		}

		for (std::size_t i = 0; i < inputs_available_.values.size(); ++i) {
			if (!inputs_available_.values[i]) {
				input_slider_state_[i] = 0;
			}
		}
	}

	void CustomHidControllerHandler::CollectInputStates()
	{
		input_buttons_state_ = {};
		input_buttons_tapped_ = {};
		input_buttons_released_ = {};
		input_slider_updates_.clear();

		for (auto& [handle, device] : devices_) {
			if (device.inputs_available_.has_any_button) {
				for (std::size_t i = 0; i < input_buttons_state_.size(); ++i) {
					if (device.inputs_available_.buttons[i]) {
						input_buttons_state_[i] |= device.input_buttons_state_[i];
						input_buttons_tapped_[i] |= device.input_buttons_tapped_[i];
						input_buttons_released_[i] |= device.input_buttons_released_[i];
					}
				}
			}
		}
		
		if (!input_slider_reports_.empty()) {
			std::array<bool, num_input_slider_values> segments_updated{};

			auto IntersectsWithUpdatedSegments = [&](decltype(SliderValuesReport::has_values) has_values) {
				for (std::size_t i = 0; i < segments_updated.size(); ++i) {
					if (segments_updated[i] && has_values[i]) {
						return true;
					}
				}
				return false;
			};

			for (const auto& update : input_slider_reports_) {
				if (IntersectsWithUpdatedSegments(update.has_values)) {
					input_slider_updates_.push_back(input_slider_state_);
					segments_updated = {};
				}
				for (std::size_t i = 0; i < input_slider_state_.size(); ++i) {
					if (!update.has_values[i]) continue;
					input_slider_state_[i] = update.values[i];
					segments_updated[i] = true;
				}
			}
			input_slider_updates_.push_back(input_slider_state_);
		}
	}

}